import pyautogui
import tkinter as tk
from PIL import ImageGrab
import colorsys
import math  # 新增数学库用于计算角度
from collections import deque  # 用于历史颜色记录


class ColorDetector:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("HSV智能取色器-环状区域采样版")

        # 屏幕中心和采样参数
        self.screen_width, self.screen_height = pyautogui.size()
        self.center_x = self.screen_width // 2
        self.center_y = self.screen_height // 2

        # 环形采样参数
        self.ring_radius = 8  # 采样环形区域半径(像素)
        self.sample_points = 8  # 采样点数
        self.dynamic_tolerance = True  # 启用动态容差

        # 历史颜色缓存
        self.color_history = deque(maxlen=10)

        # 状态变量
        self.target_hsv = None  # 改为存储HSV值
        self.h_tolerance = 15  # 色相容差（单位：度）
        self.s_tolerance = 30  # 饱和度容差（单位：%）
        self.v_tolerance = 40  # 明度容差（单位：%）
        self.detecting = False

        # 添加快捷键绑定
        self.root.bind('<Control-s>', lambda event: self.set_target_color())
        self.root.bind('<Control-S>', lambda event: self.set_target_color())

        # GUI布局
        self.setup_ui()

    def setup_ui(self):
        # 状态显示
        self.status_label = tk.Label(self.root,
                                     text="未选择目标颜色\n(使用 Ctrl+S 快速锁定)",
                                     font=('Arial', 12))
        self.status_label.pack(pady=10)

        # 当前颜色显示
        self.current_frame = tk.Frame(self.root, width=100, height=100, relief="groove", bg='white')
        self.current_frame.pack(pady=5)

        # 目标颜色显示
        self.target_frame = tk.Frame(self.root, width=100, height=100, relief="groove", bg='white')
        self.target_frame.pack(pady=5)

        # 采样参数控制面板
        param_frame = tk.Frame(self.root)
        param_frame.pack(fill=tk.X, padx=20, pady=5)

        # 采样半径设置
        tk.Label(param_frame, text="采样半径:").grid(row=0, column=0, sticky='e')
        self.radius_entry = tk.Entry(param_frame, width=5)
        self.radius_entry.insert(0, str(self.ring_radius))
        self.radius_entry.grid(row=0, column=1, sticky='w')

        # 采样点数设置
        tk.Label(param_frame, text="采样点数:").grid(row=0, column=2, sticky='e')
        self.points_entry = tk.Entry(param_frame, width=5)
        self.points_entry.insert(0, str(self.sample_points))
        self.points_entry.grid(row=0, column=3, sticky='w')

        # 容差参数控制面板
        control_frame = tk.Frame(self.root)
        control_frame.pack(fill=tk.X, padx=20, pady=5)

        # 色相容差设置
        tk.Label(control_frame, text="色相(°):").grid(row=0, column=0)
        self.h_tol_entry = tk.Entry(control_frame, width=5)
        self.h_tol_entry.insert(0, str(self.h_tolerance))
        self.h_tol_entry.grid(row=0, column=1, padx=5)

        # 饱和度容差设置
        tk.Label(control_frame, text="饱和度(%):").grid(row=0, column=2)
        self.s_tol_entry = tk.Entry(control_frame, width=5)
        self.s_tol_entry.insert(0, str(self.s_tolerance))
        self.s_tol_entry.grid(row=0, column=3, padx=5)

        # 明度容差设置
        tk.Label(control_frame, text="明度(%):").grid(row=0, column=4)
        self.v_tol_entry = tk.Entry(control_frame, width=5)
        self.v_tol_entry.insert(0, str(self.v_tolerance))
        self.v_tol_entry.grid(row=0, column=5, padx=5)

        # 动态容差选项
        self.dynamic_var = tk.BooleanVar(value=self.dynamic_tolerance)
        self.dynamic_check = tk.Checkbutton(control_frame, text="动态容差",
                                            variable=self.dynamic_var,
                                            command=self.toggle_dynamic_tolerance)
        self.dynamic_check.grid(row=0, column=6, padx=10)

        # 动态容差状态标签
        self.dynamic_status = tk.Label(control_frame, text="动态容差: 未启用", fg="gray")
        self.dynamic_status.grid(row=0, column=7)

        # 功能按钮区域
        button_frame = tk.Frame(self.root)
        button_frame.pack(pady=10)

        tk.Button(button_frame, text="锁定颜色 (Ctrl+S)",
                  command=self.set_target_color).pack(side=tk.LEFT, padx=5)
        tk.Button(button_frame, text="开始检测",
                  command=self.toggle_detection).pack(side=tk.LEFT, padx=5)
        tk.Button(button_frame, text="停止检测",
                  command=self.stop_detection).pack(side=tk.LEFT, padx=5)

        # 提示信息
        self.info_label = tk.Label(self.root,
                                   text="移动鼠标查看颜色 | Ctrl+S 锁定颜色 | Ctrl+C 退出",
                                   font=('Arial', 10))
        self.info_label.pack(pady=5)

        # 环状采样点显示（可视化采样点）
        self.canvas_frame = tk.Frame(self.root)
        self.canvas_frame.pack(pady=5)
        self.canvas = tk.Canvas(self.canvas_frame, width=100, height=100)
        self.canvas.pack()

        # 绘制初始环形采样点
        self.draw_sample_points()

    def toggle_dynamic_tolerance(self):
        self.dynamic_tolerance = self.dynamic_var.get()
        status = "启用" if self.dynamic_tolerance else "禁用"
        color = "green" if self.dynamic_tolerance else "gray"
        self.dynamic_status.config(text=f"动态容差: {status}", fg=color)

    def stop_detection(self):
        self.detecting = False
        self.status_label.config(text="检测已停止")
        self.info_label.config(text="检测已停止 | Ctrl+S 锁定颜色 | Ctrl+C 退出")
        self.root.config(bg="SystemButtonFace")
        self.canvas.delete("ring")  # 清除检测状态
        self.canvas.itemconfig("sample_point", fill="gray")  # 恢复采样点颜色

    def get_sample_points(self):
        """获取环形采样点坐标"""
        self.ring_radius = int(self.radius_entry.get() or self.ring_radius)
        self.sample_points = int(self.points_entry.get() or self.sample_points)

        points = []
        for i in range(self.sample_points):
            angle = 2 * math.pi * i / self.sample_points
            x = self.center_x + self.ring_radius * math.cos(angle)
            y = self.center_y + self.ring_radius * math.sin(angle)
            points.append((int(x), int(y)))
        return points

    def draw_sample_points(self):
        """在画布上绘制环形采样点"""
        self.canvas.delete("all")
        center_x, center_y = 50, 50  # 画布中心
        scaled_radius = min(40, self.ring_radius * 3)  # 按比例缩放

        # 绘制环形
        self.canvas.create_oval(center_x - scaled_radius, center_y - scaled_radius,
                                center_x + scaled_radius, center_y + scaled_radius,
                                outline="gray", tags="ring")

        # 绘制采样点
        for i in range(self.sample_points):
            angle = 2 * math.pi * i / self.sample_points
            x = center_x + scaled_radius * math.cos(angle)
            y = center_y + scaled_radius * math.sin(angle)
            self.canvas.create_oval(x - 3, y - 3, x + 3, y + 3,
                                    fill="gray", outline="", tags=("sample_point", f"point_{i}"))

    def get_average_color_ring(self, points):
        """获取多个采样点的平均颜色（RGB）"""
        colors = []
        try:
            # 只截取包含所有点的矩形区域
            min_x = min(p[0] for p in points)
            max_x = max(p[0] for p in points)
            min_y = min(p[1] for p in points)
            max_y = max(p[1] for p in points)

            # 截取包含所有采样点的区域
            bbox = (min_x, min_y, max_x + 1, max_y + 1)
            region = ImageGrab.grab(bbox=bbox)

            for point in points:
                # 转换为区域内的相对坐标
                rel_x = point[0] - min_x
                rel_y = point[1] - min_y

                # 确保坐标在图像范围内
                if 0 <= rel_x < region.width and 0 <= rel_y < region.height:
                    pixel = region.getpixel((rel_x, rel_y))
                    colors.append(pixel)

            # 计算平均颜色
            if colors:
                total_r, total_g, total_b = 0, 0, 0
                for color in colors:
                    total_r += color[0]
                    total_g += color[1]
                    total_b += color[2]
                n = len(colors)
                return (
                    int(total_r / n),
                    int(total_g / n),
                    int(total_b / n)
                )
            return (0, 0, 0)  # 如果没有颜色数据
        except Exception as e:
            print(f"采样错误: {e}")
            return (0, 0, 0)

    def rgb_to_hsv(self, rgb):
        """将0-255的RGB转换为0-1范围的HSV"""
        return colorsys.rgb_to_hsv(rgb[0] / 255.0, rgb[1] / 255.0, rgb[2] / 255.0)

    def hsv_distance(self, hsv1, hsv2):
        """计算HSV颜色空间的欧几里得距离（加权）"""
        # 色相分量使用圆周距离
        h_diff = min(abs(hsv1[0] - hsv2[0]), 1 - abs(hsv1[0] - hsv2[0])) * 360
        s_diff = abs(hsv1[1] - hsv2[1]) * 100
        v_diff = abs(hsv1[2] - hsv2[2]) * 100

        # 加权距离（色相权重最高）
        return math.sqrt(h_diff ** 2 * 0.6 + s_diff ** 2 * 0.3 + v_diff ** 2 * 0.1)

    def is_color_in_range(self, current_hsv):
        """HSV颜色匹配核心逻辑"""
        if not self.target_hsv:
            return False

        # 计算当前HSV与目标HSV的差异
        h_diff = min(abs(self.target_hsv[0] - current_hsv[0]),
                     1 - abs(self.target_hsv[0] - current_hsv[0])) * 360
        s_diff = abs(self.target_hsv[1] - current_hsv[1]) * 100
        v_diff = abs(self.target_hsv[2] - current_hsv[2]) * 100

        # 三重条件判断
        return (h_diff <= self.h_tolerance and
                s_diff <= self.s_tolerance and
                v_diff <= self.v_tolerance)

    def update_tolerance(self, current_hsv):
        """动态调整容差阈值（示例）"""
        if not self.color_history or len(self.color_history) < 5:
            return

        # 计算最近的HSV平均差异
        total_diff = 0
        for hist_hsv in self.color_history:
            total_diff += self.hsv_distance(hist_hsv, current_hsv)
        avg_diff = total_diff / len(self.color_history)

        # 如果环境变化大，等比例放大三个容差
        if avg_diff > 20:  # 差异阈值
            self.h_tolerance = min(30, int(self.h_tolerance * 1.2))
            self.s_tolerance = min(50, int(self.s_tolerance * 1.2))
            self.v_tolerance = min(60, int(self.v_tolerance * 1.2))

    def set_target_color(self, event=None):
        """设置目标颜色（使用环形采样）"""
        try:
            # 获取环形采样点
            points = self.get_sample_points()

            # 获取环形采样点的平均颜色
            avg_rgb = self.get_average_color_ring(points)
            if avg_rgb == (0, 0, 0):
                self.status_label.config(text="采样失败！请调整采样参数")
                return

            # 更新目标颜色和容差设置
            self.target_hsv = self.rgb_to_hsv(avg_rgb)
            self.h_tolerance = int(self.h_tol_entry.get() or 15)
            self.s_tolerance = int(self.s_tol_entry.get() or 30)
            self.v_tolerance = int(self.v_tol_entry.get() or 40)

            # 更新UI显示
            hex_color = '#{:02x}{:02x}{:02x}'.format(*avg_rgb)
            self.target_frame.config(bg=hex_color)

            status_text = f"目标颜色已锁定（环形采样）\n"
            status_text += f"半径: {self.ring_radius}px | 点数: {self.sample_points}\n"
            status_text += f"容差: H±{self.h_tolerance}° S±{self.s_tolerance}% V±{self.v_tolerance}%"
            self.status_label.config(text=status_text)

            # 绘制采样点（可视化）
            self.draw_sample_points()

            # 绿色反馈效果
            self.current_frame.config(bg='#4CAF50')
            self.root.after(200, lambda: self.current_frame.config(bg=hex_color))
        except Exception as e:
            self.status_label.config(text=f"锁定失败: {str(e)}")

    def toggle_detection(self):
        """切换检测状态"""
        if not self.target_hsv:
            self.status_label.config(text="请先锁定目标颜色！")
            return

        self.detecting = not self.detecting
        if self.detecting:
            self.status_label.config(text="检测运行中")
            self.info_label.config(text="检测运行中 | Ctrl+C 退出")
            self.color_history.clear()  # 开始检测时清空历史记录
        else:
            self.stop_detection()

    def update_color(self):
        try:
            if not self.detecting:
                self.root.after(50, self.update_color)
                return

            # 获取环形采样点
            points = self.get_sample_points()

            # 获取当前环形区域的平均颜色
            current_rgb = self.get_average_color_ring(points)
            hex_color = '#{:02x}{:02x}{:02x}'.format(*current_rgb)
            self.current_frame.config(bg=hex_color)

            # 转换为HSV用于检测
            current_hsv = self.rgb_to_hsv(current_rgb)

            # 添加到历史记录
            self.color_history.append(current_hsv)

            # 动态调整容差
            if self.dynamic_tolerance:
                self.update_tolerance(current_hsv)

            # 颜色检测逻辑
            if self.detecting and self.target_hsv:
                # 更新画布显示
                self.canvas.itemconfig("ring", outline="#4CAF50" if self.detecting else "gray")

                # 检查是否匹配
                match_found = self.is_color_in_range(current_hsv)

                # 更新采样点颜色状态
                for i in range(self.sample_points):
                    if match_found:
                        self.canvas.itemconfig(f"point_{i}", fill="#FF5733")  # 匹配时红色
                    else:
                        self.canvas.itemconfig(f"point_{i}", fill="#3498DB")  # 未匹配时蓝色

                # 更新信息显示
                if match_found:
                    self.root.config(bg="#4CAF50")
                    self.info_label.config(text="★ 检测到目标颜色！ ★")
                else:
                    # 计算最小差异值
                    h_diff = min(abs(self.target_hsv[0] - current_hsv[0]),
                                 1 - abs(self.target_hsv[0] - current_hsv[0])) * 360
                    s_diff = abs(self.target_hsv[1] - current_hsv[1]) * 100
                    v_diff = abs(self.target_hsv[2] - current_hsv[2]) * 100

                    # 格式化为显示文本
                    diff_text = f"H:{h_diff:.1f}° S:{s_diff:.1f}% V:{v_diff:.1f}%"
                    tol_text = f"(容差: H±{self.h_tolerance}° S±{self.s_tolerance}% V±{self.v_tolerance}%)"

                    self.root.config(bg="SystemButtonFace")
                    self.info_label.config(text=f"检测中... | 当前差异: {diff_text} {tol_text}")

            self.root.after(50, self.update_color)
        except Exception as e:
            print(f"更新错误: {e}")
            self.root.after(50, self.update_color)

    def run(self):
        self.root.after(0, self.update_color)
        self.root.mainloop()


if __name__ == "__main__":
    detector = ColorDetector()
    try:
        detector.run()
    except KeyboardInterrupt:
        print("程序已安全退出")